import argparse
import os
import sys
import hashlib
from time import sleep
import grpc
import threading
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# 多播组模态绑定（与需求对齐，增加注释说明）
MCAST_GROUP_MODES = {
    1: {
        'modes': ['ipv4', 'ipv6'], 
        'mac': "01:00:5e:00:00:01",
        'description': "IPv4 + IPv6 多播组"
    },
    2: {
        'modes': ['source_route', 'ipv4'], 
        'mac': "01:00:5e:00:00:02",
        'description': "源路由 + IPv4 多播组"
    },
    3: {
        'modes': ['source_route', 'ipv6'], 
        'mac': "33:33:ff:00:00:03",
        'description': "源路由 + IPv6 多播组"
    },
    4: {
        'modes': ['ipv4', 'ipv6', 'source_route'], 
        'mac': "33:33:ff:00:00:04",
        'description': "全模态多播组"
    }
}

# 模态端口映射（示例，需与拓扑一致，增加注释）
MODE_PORT_MAP = {
    'ipv4': 1,
    'ipv6': 2,
    'source_route': 3
}

# 错误处理装饰器
def handle_errors(func):
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except grpc.RpcError as e:
            printGrpcError(e)
            print(f"Error in {func.__name__} with args: {args}, {kwargs}")
            sys.exit(1)
    return wrapper

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))

@handle_errors
def writeIpv4Rules(p4info_helper, switch, dst_ip_addr, dst_eth_addr, port):
    """下发 IPv4 单播转发规则"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": (dst_ip_addr, 32)},
        action_name="MyIngress.ipv4_forward",
        action_params={"dstAddr": dst_eth_addr, "port": port})
    switch.WriteTableEntry(table_entry)
    print(f"Installed IPv4 rule on {switch.name}: {dst_ip_addr} -> port {port}")

@handle_errors
def writeArpRules(p4info_helper, switch, dst_ip_addr, dst_eth_addr, port):
    """下发 ARP 应答规则"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.arp_table",
        match_fields={"hdr.arp.tpa": dst_ip_addr},
        action_name="MyIngress.arp_reply",
        action_params={"dstAddr": dst_eth_addr, "port": port})
    switch.WriteTableEntry(table_entry)
    print(f"Installed ARP rule on {switch.name}: {dst_ip_addr} -> port {port}")

@handle_errors
def writeIpv6Rules(p4info_helper, switch, dst_ip_addr, prefix_length, dst_eth_addr, port):
    """下发 IPv6 单播转发规则"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv6_lpm",
        match_fields={"hdr.ipv6.dstAddr": (dst_ip_addr, prefix_length)},
        action_name="MyIngress.ipv6_forward",
        action_params={"dstAddr": dst_eth_addr, "port": port})
    switch.WriteTableEntry(table_entry)
    print(f"Installed IPv6 rule on {switch.name}: {dst_ip_addr}/{prefix_length} -> port {port}")

@handle_errors
def writeSourceRoutingRules(p4info_helper, switch, current_index, port):
    """下发源路由转发规则"""
    action_name = f"MyIngress.source_route_forward_{current_index}"
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.sr_table",
        match_fields={"hdr.sr.current_index": current_index},
        action_name=action_name,
        action_params={})
    switch.WriteTableEntry(table_entry)
    print(f"Installed SR rule on {switch.name}: index {current_index} -> port {port}")

@handle_errors
def write_multicast_rules(p4info_helper, switch, mcast_mac, group_id):
    """下发多播组匹配规则（关联模态掩码）"""
    # 从组 ID 映射模态掩码
    mode_mask = 0
    for mode in MCAST_GROUP_MODES[group_id]['modes']:
        if mode in MODE_PORT_MAP:
            mode_bit = list(MODE_PORT_MAP.keys()).index(mode)
            mode_mask |= (1 << mode_bit)
    
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.multicast_table",
        match_fields={"hdr.ethernet.dstAddr": mcast_mac},
        action_name="MyIngress.multicast_forward",
        action_params={
            "group_id": group_id,
            "mode_mask": mode_mask  # 传递模态掩码给 P4
        }
    )
    switch.WriteTableEntry(table_entry)
    print(f"Installed multicast rule on {switch.name}: MAC {mcast_mac} -> Group {group_id}")

@handle_errors
def setup_multicast_groups(p4info_helper, switch, group_id, ports):
    """配置多播组出端口"""
    replicas = []
    for port in ports:
        replicas.append({"egress_port": port, "instance": 0})
    mc_entry = p4info_helper.buildMulticastGroupEntry(
        multicast_group_id=group_id, 
        replicas=replicas
    )
    switch.WritePREEntry(mc_entry)
    group_info = MCAST_GROUP_MODES[group_id]
    print(f"Configured multicast group {group_id} on {switch.name} "
          f"({group_info['description']}) with ports {ports}")

@handle_errors
def writeModeSelectionRules(p4info_helper, switch, src_ip, dst_ip, num_modes, base_port, src_ipv6=None, dst_ipv6=None):
    """下发模态调度规则（S1/S4 专用，支持IPv4和IPv6）"""
    match_fields = {
        "hdr.ipv4.srcAddr": src_ip,
        "hdr.ipv4.dstAddr": dst_ip
    }
    # 若提供了IPv6地址，添加IPv6匹配字段
    if src_ipv6 and dst_ipv6:
        match_fields.update({
            "hdr.ipv6.srcAddr": src_ipv6,
            "hdr.ipv6.dstAddr": dst_ipv6
        })
    
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.mode_selection_table",
        match_fields=match_fields,
        action_name="MyIngress.mode_schedule",
        action_params={
            "num_modes": num_modes,
            "base_port": base_port
        }
    )
    switch.WriteTableEntry(table_entry)
    if src_ipv6 and dst_ipv6:
        print(f"Installed mode selection rule on {switch.name}: "
              f"IPv4({src_ip}->{dst_ip}) & IPv6({src_ipv6}->{dst_ipv6})")
    else:
        print(f"Installed mode selection rule on {switch.name}: "
              f"IPv4({src_ip}->{dst_ip})")

def verify_switch_connection(switch):
    """验证交换机连接状态"""
    try:
        switch.MasterArbitrationUpdate()
        return True
    except grpc.RpcError:
        return False

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 初始化交换机连接（S1/S4 为汇总控制，S2/S3/S5 为运输）
        switches = {
            's1': p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name='s1', address='127.0.0.1:50051', device_id=0,
                proto_dump_file='logs/s1-p4runtime-requests.txt'),
            's2': p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name='s2', address='127.0.0.1:50052', device_id=1,
                proto_dump_file='logs/s2-p4runtime-requests.txt'),
            's3': p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name='s3', address='127.0.0.1:50053', device_id=2,
                proto_dump_file='logs/s3-p4runtime-requests.txt'),
            's4': p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name='s4', address='127.0.0.1:50054', device_id=3,
                proto_dump_file='logs/s4-p4runtime-requests.txt'),
            's5': p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name='s5', address='127.0.0.1:50055', device_id=4,
                proto_dump_file='logs/s5-p4runtime-requests.txt')
        }

        # 验证所有交换机连接
        for name, sw in switches.items():
            if not verify_switch_connection(sw):
                print(f"Failed to connect to switch {name}")
                sys.exit(1)

        # 加载 P4 程序
        for name, sw in switches.items():
            sw.SetForwardingPipelineConfig(
                p4info=p4info_helper.p4info,
                bmv2_json_file_path=bmv2_file_path)
            print(f"Installed P4 program on {name}")

        # ---------------------------
        # 1. 单播转发规则（运输交换机：S2/S3/S5；汇总交换机：S1/S4 需兼容所有模态）
        # ---------------------------
        # S2（IPv4 模态）
        writeIpv4Rules(p4info_helper, switches['s2'], "10.0.2.2", "08:00:00:00:02:22", 2)
        writeIpv4Rules(p4info_helper, switches['s2'], "10.0.1.1", "08:00:00:00:01:11", 1)
        # S3（IPv6 模态）
        writeIpv6Rules(p4info_helper, switches['s3'], "fe80::5678", 128, "08:00:00:00:02:22", 2)
        # S5（源路由模态）
        writeSourceRoutingRules(p4info_helper, switches['s5'], 1, 2)
        # S1/S4（全模态兼容）
        for sw_name in ['s1', 's4']:
            sw = switches[sw_name]
            writeIpv4Rules(p4info_helper, sw, "10.0.2.2", "08:00:00:00:02:22", 2)
            writeIpv4Rules(p4info_helper, sw, "10.0.1.1", "08:00:00:00:01:11", 1)
            writeIpv6Rules(p4info_helper, sw, "fe80::5678", 128, "08:00:00:00:02:22", 3)
            writeSourceRoutingRules(p4info_helper, sw, 0, 4)  # 源路由起始索引

        # ---------------------------
        # 2. ARP 规则（仅 S1/S4 需处理多模态终端）
        # ---------------------------
        writeArpRules(p4info_helper, switches['s1'], "10.0.1.1", "08:00:00:00:01:11", 1)
        writeArpRules(p4info_helper, switches['s1'], "10.0.1.10", "08:00:00:00:01:00", 1)
        writeArpRules(p4info_helper, switches['s4'], "10.0.2.2", "08:00:00:00:02:22", 3)
        writeArpRules(p4info_helper, switches['s4'], "10.0.2.20", "08:00:00:00:02:00", 3)

        # ---------------------------
        # 3. 多播组规则（仅 S1/S4 需配置）
        # ---------------------------
        for group_id in MCAST_GROUP_MODES:
            mcast_mac = MCAST_GROUP_MODES[group_id]['mac']
            # 下发多播组匹配规则（关联模态掩码）
            write_multicast_rules(p4info_helper, switches['s1'], mcast_mac, group_id)
            write_multicast_rules(p4info_helper, switches['s4'], mcast_mac, group_id)
            # 配置多播组出端口（示例：S1 出端口 2/3/4；S4 出端口 1/3/4）
            if group_id == 1:
                setup_multicast_groups(p4info_helper, switches['s1'], group_id, [2, 3, 4])
                setup_multicast_groups(p4info_helper, switches['s4'], group_id, [1, 3, 4])
            else:
                setup_multicast_groups(p4info_helper, switches['s1'], group_id, [2, 4])
                setup_multicast_groups(p4info_helper, switches['s4'], group_id, [1, 3])

        # ---------------------------
        # 4. 模态调度规则（S1/S4 专用，实现多播组动态模态选择）
        # ---------------------------
        # IPv4 模态调度
        writeModeSelectionRules(p4info_helper, switches['s1'], "10.0.1.1", "10.0.2.2", 3, 1)
        writeModeSelectionRules(p4info_helper, switches['s4'], "10.0.2.2", "10.0.1.1", 3, 1)
        # IPv6 模态调度（示例）
        ipv6_src = "2001:db8::1"
        ipv6_dst = "2001:db8::2"
        writeModeSelectionRules(p4info_helper, switches['s1'], "10.0.1.1", "10.0.2.2", 3, 1, ipv6_src, ipv6_dst)
        writeModeSelectionRules(p4info_helper, switches['s4'], "10.0.2.2", "10.0.1.1", 3, 1, ipv6_dst, ipv6_src)

        print("\nAll rules installed. Press Ctrl+C to stop.")
        while True:
            sleep(1)

    except KeyboardInterrupt:
        print("\nShutting down.")
    except Exception as e:
        print(f"Unexpected error: {e}")
    finally:
        ShutdownAllSwitchConnections()
        print("All switch connections closed.")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller with enhanced error handling')
    parser.add_argument('--p4info', help='p4info proto in text format',
                        type=str, default='./build/basic.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file',
                        type=str, default='./build/basic.json')
    args = parser.parse_args()

    # 检查文件存在性
    if not os.path.exists(args.p4info):
        print(f"P4Info file not found: {args.p4info}")
        sys.exit(1)
    if not os.path.exists(args.bmv2_json):
        print(f"BMv2 JSON file not found: {args.bmv2_json}")
        sys.exit(1)

    main(args.p4info, args.bmv2_json)